#include "partitioner.h"
#include "display.h"
#include "esp_heap_caps.h"
#include "mykeyboard.h"
#include "sd_functions.h"
#include <globals.h>

// Define o tamanho da partição
#define PARTITION_SIZE 4096

// Using "buff[4096]" to store and write the partitions
#if defined(PART_08MB)
// default partition scheme(App, FAT and SPIFFS)
const uint8_t def_part[224] PROGMEM = {
    0xAA, 0x50, 0x01, 0x02, 0x00, 0x90, 0x00, 0x00, 0x00, 0x60, 0x00, 0x00, 0x6E, 0x76, 0x73, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xAA, 0x50, 0x00, 0x20, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x16, 0x00, 0x61, 0x70, 0x70, 0x30,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xAA, 0x50, 0x00, 0x10, 0x00, 0x00, 0x17, 0x00, 0x00, 0x00, 0x4F, 0x00, 0x61, 0x70, 0x70, 0x31,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xAA, 0x50, 0x01, 0x81, 0x00, 0x00, 0x67, 0x00, 0x00, 0x00, 0x08, 0x00, 0x76, 0x66, 0x73, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xAA, 0x50, 0x01, 0x82, 0x00, 0x00, 0x6F, 0x00, 0x00, 0x00, 0x10, 0x00, 0x73, 0x70, 0x69, 0x66,
    0x66, 0x73, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xAA, 0x50, 0x01, 0x03, 0x00, 0x00, 0x7F, 0x00, 0x00, 0x00, 0x01, 0x00, 0x63, 0x6F, 0x72, 0x65,
    0x64, 0x75, 0x6D, 0x70, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xEB, 0xEB, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
    0xBF, 0xE1, 0xC0, 0x6C, 0x4F, 0xCC, 0x25, 0x52, 0x53, 0xAB, 0xDA, 0xEA, 0x74, 0x87, 0x7F, 0x13
};
// 6Mb app partition
const uint8_t doom[160] PROGMEM = {
    0xAA, 0x50, 0x01, 0x02, 0x00, 0x90, 0x00, 0x00, 0x00, 0x60, 0x00, 0x00, 0x6E, 0x76, 0x73, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xAA, 0x50, 0x00, 0x20, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x16, 0x00, 0x61, 0x70, 0x70, 0x30,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xAA, 0x50, 0x00, 0x10, 0x00, 0x00, 0x17, 0x00, 0x00, 0x00, 0x68, 0x00, 0x61, 0x70, 0x70, 0x31,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xAA, 0x50, 0x01, 0x03, 0x00, 0x00, 0x7F, 0x00, 0x00, 0x00, 0x01, 0x00, 0x63, 0x6F, 0x72, 0x65,
    0x64, 0x75, 0x6D, 0x70, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xEB, 0xEB, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
    0x2D, 0x3C, 0x30, 0x3F, 0x42, 0x7D, 0x51, 0xEE, 0xE9, 0x2B, 0x8D, 0x78, 0xCB, 0x29, 0x7D, 0xDC
};

const uint8_t uiflow2[192] PROGMEM = {
    // uiflow partition scheme, APP, sys(FAT) and vfs(FAT)
    0xAA, 0x50, 0x01, 0x02, 0x00, 0x90, 0x00, 0x00, 0x00, 0x60, 0x00, 0x00, 0x6E, 0x76, 0x73, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xAA, 0x50, 0x00, 0x20, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x16, 0x00, 0x61, 0x70, 0x70, 0x30,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xAA, 0x50, 0x00, 0x10, 0x00, 0x00, 0x17, 0x00, 0x00, 0x00, 0x4E, 0x00, 0x61, 0x70, 0x70, 0x31,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xAA, 0x50, 0x01, 0x81, 0x00, 0x00, 0x65, 0x00, 0x00, 0x00, 0x10, 0x00, 0x73, 0x79, 0x73, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xAA, 0x50, 0x01, 0x81, 0x00, 0x00, 0x75, 0x00, 0x00, 0x00, 0x0B, 0x00, 0x76, 0x66, 0x73, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xEB, 0xEB, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
    0xF0, 0x4F, 0xA2, 0x1D, 0x91, 0x76, 0x30, 0x87, 0x76, 0x59, 0xCC, 0x84, 0xED, 0x69, 0x02, 0xE3
};

const uint8_t gamestation[192] PROGMEM = {
    // PArtition scheme for Cardputer Gamestation with 4,5Mb of SPIFFS
    0xAA, 0x50, 0x01, 0x02, 0x00, 0x90, 0x00, 0x00, 0x00, 0x60, 0x00, 0x00, 0x6E, 0x76, 0x73, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xAA, 0x50, 0x00, 0x20, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x16, 0x00, 0x61, 0x70, 0x70, 0x30,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xAA, 0x50, 0x00, 0x10, 0x00, 0x00, 0x17, 0x00, 0x00, 0x00, 0x20, 0x00, 0x61, 0x70, 0x70, 0x31,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xAA, 0x50, 0x01, 0x82, 0x00, 0x00, 0x37, 0x00, 0x00, 0x00, 0x48, 0x00, 0x73, 0x70, 0x69, 0x66,
    0x66, 0x73, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xAA, 0x50, 0x01, 0x03, 0x00, 0x00, 0x7F, 0x00, 0x00, 0x00, 0x01, 0x00, 0x63, 0x6F, 0x72, 0x65,
    0x64, 0x75, 0x6D, 0x70, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xEB, 0xEB, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
    0x48, 0xF3, 0xAB, 0xB2, 0x54, 0x6D, 0xCB, 0x5E, 0xF2, 0x78, 0x96, 0xD4, 0xF6, 0x64, 0x3C, 0x1F
};
#elif defined(PART_16MB)

const uint8_t def_part[288] PROGMEM = {
    0xAA, 0x50, 0x01, 0x02, 0x00, 0x90, 0x00, 0x00, 0x00, 0x60, 0x00, 0x00, 0x6E, 0x76, 0x73, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xAA, 0x50, 0x00, 0x20, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x1F, 0x00, 0x61, 0x70, 0x70, 0x30,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xAA, 0x50, 0x00, 0x10, 0x00, 0x00, 0x20, 0x00, 0x00, 0x00, 0x80, 0x00, 0x61, 0x70, 0x70, 0x31,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xAA, 0x50, 0x01, 0x81, 0x00, 0x00, 0xA0, 0x00, 0x00, 0x00, 0x10, 0x00, 0x73, 0x79, 0x73, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xAA, 0x50, 0x01, 0x81, 0x00, 0x00, 0xB0, 0x00, 0x00, 0x00, 0x20, 0x00, 0x76, 0x66, 0x73, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xAA, 0x50, 0x01, 0x82, 0x00, 0x00, 0xD0, 0x00, 0x00, 0x00, 0x2F, 0x00, 0x73, 0x70, 0x69, 0x66,
    0x66, 0x73, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xAA, 0x50, 0x01, 0x03, 0x00, 0x00, 0xFF, 0x00, 0x00, 0x00, 0x01, 0x00, 0x63, 0x6F, 0x72, 0x65,
    0x64, 0x75, 0x6D, 0x70, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xEB, 0xEB, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
    0x2C, 0x4E, 0x70, 0x13, 0x8D, 0xF3, 0xB0, 0xF7, 0xBF, 0x69, 0x7C, 0xF1, 0x13, 0xDB, 0x36, 0xC1
};

const uint8_t uiFlow1[352] PROGMEM = {
    0xAA, 0x50, 0x01, 0x02, 0x00, 0x90, 0x00, 0x00, 0x00, 0x60, 0x00, 0x00, 0x6E, 0x76, 0x73, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xAA, 0x50,
    0x01, 0x01, 0x00, 0xF0, 0x00, 0x00, 0x00, 0x10, 0x00, 0x00, 0x70, 0x68, 0x79, 0x5F, 0x69, 0x6E, 0x69,
    0x74, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xAA, 0x50, 0x00, 0x20,
    0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x1D, 0x00, 0x61, 0x70, 0x70, 0x30, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xAA, 0x50, 0x00, 0x10, 0x00, 0x00,
    0x1E, 0x00, 0x00, 0x00, 0x4E, 0x00, 0x61, 0x70, 0x70, 0x31, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xAA, 0x50, 0x01, 0x81, 0x00, 0x00, 0x6C, 0x00,
    0x00, 0x00, 0x53, 0x00, 0x76, 0x66, 0x73, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xAA, 0x50, 0x40, 0x40, 0x00, 0xE0, 0xBF, 0x00, 0x00, 0x10,
    0x00, 0x00, 0x63, 0x6F, 0x6E, 0x66, 0x69, 0x67, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0xAA, 0x50, 0x50, 0x50, 0x00, 0xF0, 0xBF, 0x00, 0x00, 0x10, 0x00, 0x00,
    0x77, 0x69, 0x66, 0x69, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0xAA, 0x50, 0x01, 0x81, 0x00, 0x00, 0xC0, 0x00, 0x00, 0x00, 0x10, 0x00, 0x73, 0x79,
    0x73, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0xAA, 0x50, 0x01, 0x82, 0x00, 0x00, 0xD0, 0x00, 0x00, 0x00, 0x2F, 0x00, 0x73, 0x70, 0x69, 0x66,
    0x66, 0x73, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xAA,
    0x50, 0x01, 0x03, 0x00, 0x00, 0xFF, 0x00, 0x00, 0x00, 0x01, 0x00, 0x63, 0x6F, 0x72, 0x65, 0x64, 0x75,
    0x6D, 0x70, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xEB, 0xEB, 0xFF,
    0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x38, 0x4C, 0x68, 0xD6,
    0x6A, 0x40, 0x6E, 0x11, 0xB8, 0x86, 0xC8, 0xA7, 0xBE, 0xD5, 0x72, 0xF9
};
#endif

// Função para apagar e escrever na região de memória flash
bool partitionSetter(const uint8_t *scheme, size_t scheme_size) {
    uint8_t *buffer = (uint8_t *)heap_caps_malloc(PARTITION_SIZE, MALLOC_CAP_INTERNAL);
    if (buffer == NULL) {
        ESP_LOGE("FLASH", "Failed to allocate buffer in DRAM");
        return false;
    }

    // Preencher o buffer com 0xFF
    memset(buffer, 0xFF, PARTITION_SIZE);

    // Copiar o esquema de partição para o buffer
    memcpy(buffer, scheme, scheme_size);

    esp_err_t err;

    // Apagar a região de memória flash
    err = esp_flash_erase_region(NULL, 0x8000, PARTITION_SIZE);
    if (err != ESP_OK) {
        ESP_LOGE("FLASH", "Failed to erase flash region (0x%x)", err);
        heap_caps_free(buffer);
        return false;
    }

    // Escrever o buffer na memória flash
    err = esp_flash_write(NULL, buffer, 0x8000, PARTITION_SIZE);
    if (err != ESP_OK) {
        ESP_LOGE("FLASH", "Failed to write to flash (0x%x)", err);
        heap_caps_free(buffer);
        return false;
    }

    heap_caps_free(buffer);
    return true;
}

void partitioner() {
#if CONFIG_IDF_TARGET_ESP32P4
    const esp_partition_t *part =
        esp_partition_find_first(ESP_PARTITION_TYPE_APP, ESP_PARTITION_SUBTYPE_APP_OTA_0, NULL);
#endif
    int partition = 100;
    const uint8_t *data = nullptr;
    size_t data_size = 0;

    // Opções de partição
    options = {
        {"Default",      [&]() { partition = 0; }},
#if defined(PART_08MB)
        {"Doom",         [&]() { partition = 1; }},
        {"UiFlow2",      [&]() { partition = 2; }},
        {"Game Station", [&]() { partition = 3; }},
#elif defined(PART_04MB)
        {"Orca", [&]() { partition = 1; }},
#elif defined(PART_16MB)
        {"UiFlow1", [&]() { partition = 1; }},
#endif
    };
    loopOptions(options);

    if (partition == 100) goto Exit;
    switch (partition) {
#if !defined(PART_04MB)
        case 0:
            data = def_part;
            data_size = sizeof(def_part);
            break;
#endif
#if defined(PART_08MB)
        case 1:
            data = doom;
            data_size = sizeof(doom);
            break;
        case 2:
            data = uiflow2;
            data_size = sizeof(uiflow2);
            break;
        case 3:
            data = gamestation;
            data_size = sizeof(gamestation);
            break;
#elif defined(PART_16MB)
        case 1:
            data = uiFlow1;
            displayRedStripe("Experimental");
            delay(2500);
            data_size = sizeof(uiFlow1);
            break;
#endif
        default: goto Exit;
    }

    if (!partitionSetter(data, data_size)) {
        Serial.println("Error when running partitionSetter function");
        displayRedStripe("Partitioning Error");
        while (!check(SelPress)) yield();
    }

    displayRedStripe("Restart needed");

    while (!check(SelPress)) yield();
    while (check(SelPress)) yield();
    FREE_TFT
#if CONFIG_IDF_TARGET_ESP32P4
    esp_ota_set_boot_partition(part);
    ESP.deepSleep(100);
#endif
    ESP.restart();
Exit:
    Serial.print("Desistiu");
}

void partList() {
    // Obtemos a lista de partições
    const esp_partition_t *partition;
    esp_partition_iterator_t it = esp_partition_find(ESP_PARTITION_TYPE_ANY, ESP_PARTITION_SUBTYPE_ANY, NULL);

    if (it != NULL) {
        Serial.println("Partições encontradas:");
        String txt = "";
        int i = 0;
        while (it != NULL) {
            partition = esp_partition_get(it);

            switch (partition->subtype) {
                case ESP_PARTITION_SUBTYPE_APP_OTA_0:
                case ESP_PARTITION_SUBTYPE_APP_OTA_1:
                    Serial.println("OTA");
                    txt += "-OTA-";
                    break;

                case ESP_PARTITION_SUBTYPE_DATA_FAT:
                    Serial.println("FAT");
                    txt += "FAT-";
                    break;
                case ESP_PARTITION_SUBTYPE_DATA_SPIFFS:
                    Serial.println("SPIFFS");
                    txt += "SPIFFs-";
                    break;
                default: Serial.println("Desconhecido"); break;
            }
            it = esp_partition_next(it);
        }
        esp_partition_iterator_release(it);
        displayRedStripe(txt);
        delay(300);
        while (!check(SelPress)) yield();
        while (check(SelPress)) yield();
    }
}

void dumpPartition(const char *partitionLabel, const char *outputPath) {
    tft->fillScreen(BGCOLOR);
    const esp_partition_t *partition =
        esp_partition_find_first(ESP_PARTITION_TYPE_DATA, ESP_PARTITION_SUBTYPE_ANY, partitionLabel);
    if (partition == NULL) {
        Serial.printf("Partição %s não encontrada\n", partitionLabel);
        return;
    }

    setupSdCard();
    if (!SDM.exists("/bkp")) SDM.mkdir("/bkp");

    String output = outputPath;
    output += ".bin";
    int i = 0;
    while (SDM.exists(output)) {
        i++;
        output = String(outputPath) + String(i) + ".bin";
    }

    File outputFile = SDM.open(output.c_str(), FILE_WRITE, true);
    if (!outputFile) {
        Serial.printf("Falha ao abrir o arquivo %s no cartão SD\n", outputPath);
        return;
    }

    Serial.printf("Iniciando dump da partição %s para o arquivo %s\n", partitionLabel, outputPath);

    const size_t bufferSize = 1024; // Ajuste conforme necessário
    uint8_t buffer[1024];
    size_t bytesToRead = 0;
    esp_err_t result;
    progressHandler(0, 500);
    displayRedStripe("Backing up");
    for (size_t offset = 0; offset < partition->size; offset += bufferSize) {
        bytesToRead = (offset + bufferSize > partition->size) ? (partition->size - offset) : bufferSize;
        result = esp_partition_read(partition, offset, buffer, bytesToRead);
        if (result != ESP_OK) {
            Serial.printf(
                "Erro ao ler a partição %s no offset %d (código de erro: %d)\n",
                partitionLabel,
                offset,
                result
            );
            outputFile.close();
            return;
        }
        outputFile.write(buffer, bytesToRead);
        progressHandler(int(offset + bufferSize), partition->size);
    }
    outputFile.close();
    displayRedStripe("    Complete!    ");
    delay(500);
    Serial.printf("Dump da partição %s para o arquivo %s concluído\n", partitionLabel, outputPath);

    bool attach = false;
    options = {
        {"Attach to a file", [&]() { attach = true; }      },
        {"Main Menu",        [=]() { returnToMenu = true; }}
    };
    loopOptions(options);

    if (attach) {
        String to = loopSD(true);
        attachPartition(output, to);
    }
}

void restorePartition(const char *partitionLabel) {
    String filepath = loopSD(true);
    tft->fillScreen(BGCOLOR);
    if (filepath == "") return;
    else {
        File source = SDM.open(filepath, "r");
        if (strcmp(partitionLabel, "spiffs") == 0) {
            prog_handler = 1;
            Update.begin(source.size(), U_SPIFFS);
            uint8_t buffer[1024];
            int bytesRead = 0;
            int written = 0;
            size_t total = source.size();
            progressHandler(0, 500);
            while (source.available()) {
                bytesRead = source.read(buffer, sizeof(buffer));
                Update.write(buffer, bytesRead);
                written += bytesRead;
                progressHandler(written, total);
            }
        }

        if (strcmp(partitionLabel, "vfs") == 0) { performFATUpdate(source, source.size(), "vfs"); }
        if (strcmp(partitionLabel, "sys") == 0) { performFATUpdate(source, source.size(), "sys"); }
    }
    delay(100);
    displayRedStripe("    Restored!    ");
    delay(2500);
}

#define TAG "Partitioneer"
#define BUFFER_SIZE 1024

// Função para copiar partições com buffer de 1024 bytes
esp_err_t copy_partition(const esp_partition_t *src, const esp_partition_t *dst) {
    uint8_t buffer[BUFFER_SIZE];
    esp_err_t err;
    progressHandler(0, 500);
    displayRedStripe("Launcher Update");
    for (size_t offset = 0; offset < dst->size; offset += BUFFER_SIZE) {
        size_t read_size = BUFFER_SIZE;
        if (offset + BUFFER_SIZE > dst->size) { read_size = dst->size - offset; }

        err = esp_partition_read(src, offset, buffer, read_size);
        if (err != ESP_OK) {
            ESP_LOGE(TAG, "Failed to read source partition at offset %u", offset);
            return err;
        }

        err = esp_partition_write(dst, offset, buffer, read_size);
        if (err != ESP_OK) {
            ESP_LOGE(TAG, "Failed to write to destination partition at offset %u", offset);
            return err;
        }
        progressHandler(offset + BUFFER_SIZE, dst->size);
    }

    return ESP_OK;
}
#if CONFIG_IDF_TARGET_ESP32P4
#define TARGET_PARTITION ESP_PARTITION_SUBTYPE_APP_FACTORY
#else
#define TARGET_PARTITION ESP_PARTITION_SUBTYPE_APP_TEST
#endif
// Função principal
void partitionCrawler() {
    const esp_partition_t *running_partition = esp_ota_get_running_partition();
    if (running_partition == NULL) {
        ESP_LOGE(TAG, "Failed to get running partition");
        return;
    }

    if (running_partition->subtype == TARGET_PARTITION) {
        ESP_LOGI(
            TAG,
            "Running partition is %s partition, no action taken",
            TARGET_PARTITION == ESP_PARTITION_SUBTYPE_APP_TEST ? "TEST" : "FACTORY"
        );
        return;
    }

    displayRedStripe("Updating...");

    const esp_partition_t *test_partition =
        esp_partition_find_first(ESP_PARTITION_TYPE_APP, TARGET_PARTITION, NULL);

    if (test_partition == NULL) {
        ESP_LOGE(TAG, "Failed to find test partition");
        return;
    }

    ESP_LOGI(TAG, "Erasing test partition");
    esp_err_t err = esp_partition_erase_range(test_partition, 0, test_partition->size);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "Failed to erase test partition");
        return;
    }

    ESP_LOGI(TAG, "Copying running partition to test partition");
    err = copy_partition(running_partition, test_partition);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "Failed to copy partition data");
        displayRedStripe("Use M5Burner!");
        delay(5000);
        return;
    }

    ESP_LOGI(TAG, "Writing 0x00 to first byte of the running partition (break OTA0 Launcher)");
    uint8_t zero_byte = 0x00;
    err = esp_partition_write(running_partition, 0, &zero_byte, 1);

    if (err != ESP_OK) {
        ESP_LOGE(TAG, "Failed to write 0x00 to the first byte of the running partition");
    } else {
        ESP_LOGI(TAG, "Restarting system to boot from test partition");
        esp_restart();
    }
}

bool attachPartition(String _from, String _to) {
    size_t offset = 0;
    uint8_t bytes[16];
    Serial.printf("From: %s\nTo: %s\n", _from.c_str(), _to.c_str());
    File to = SDM.open(_to, FILE_READ);
    if (!to) {
        displayRedStripe("Can't open target");
        delay(2500);
        return false;
    }

    // look fot FAT/SPIFFS/LittleFS partition offset
    to.seek(0x8000);
    to.read(bytes, 16);
    if (bytes[0] != 0xAA || bytes[1] != 0x50 || bytes[2] != 0x01) {
        displayRedStripe("Has no partition table");
        delay(2500);
        return false; // couldn't find a valid position
    }
    for (int i = 0x0; i <= 0x1A0; i += 0x20) {
        if (!to.seek(0x8000 + i)) {
            Serial.println("Error: Could not move cursor to read partition info");
            to.close();
            return false;
        }
        to.read(bytes, 16);

        if (bytes[3] == 0x81 || bytes[3] == 0x82 || bytes[3] == 0x83) {
            // Calculate offset (big endian)
            offset = (bytes[0x06] << 16) | (bytes[0x07] << 8) | bytes[0x08];
            Serial.printf("offset=%d\n", offset);
        }
    }
    to.close();

    if (!offset) {
        displayRedStripe("Invalid target");
        delay(2500);
        return false; // couldn't find a valid position
    }

    File target = SDM.open(_to, FILE_WRITE);
    if (!target) {
        displayRedStripe("Can't reopen target");
        delay(2500);
        return false;
    }

    // Adjust the target file size
    if (target.size() > offset) {
        // Erases the old partition with 0xFF
        target.seek(offset);
        while (target.position() < target.size()) {
            int angle = (360 * (target.position() - offset)) / (target.size() - offset);
            tft->drawArc(tftWidth / 2, tftHeight / 2, 50, 45, 0, angle, FGCOLOR);
            target.write(0xFF);
        }
    } else if (target.size() < offset) {
        // fill with 0xFF until offset
        target.seek(target.size());
        size_t fillLen = offset - target.size();
        const size_t chunk = 512;
        uint8_t buf[chunk];
        memset(buf, 0xFF, chunk);
        while (fillLen > 0) {
            int angle = (360 * (offset - fillLen)) / offset;
            tft->drawArc(tftWidth / 2, tftHeight / 2, 40, 35, 0, angle, FGCOLOR);
            size_t w = min(chunk, fillLen);
            target.write(buf, w);
            fillLen -= w;
        }
    }

    // copy data from source file
    File from = SDM.open(_from, FILE_READ);
    if (!from) {
        displayRedStripe("Can't open source");
        delay(2500);
        target.close();
        return false;
    }

    target.seek(offset);
    while (true) {
        int angle = (360 * (target.position() - offset)) / from.size();
        tft->drawArc(tftWidth / 2, tftHeight / 2, 35, 30, 0, angle, ALCOLOR);
        size_t bytesRead = from.read(buff, sizeof(buff));
        if (!bytesRead) break;
        target.write(buff, bytesRead);
    }

    from.close();
    target.close();

    Serial.printf(
        "Partition data from '%s' attached at offset 0x%X into '%s'\n",
        _from.c_str(),
        (unsigned int)offset,
        _to.c_str()
    );

    return true;
}
